from sqlalchemy import select
from config.get_db import get_db
import json
from module_stock.entity.do.stock_basic_do import StockBasic
from module_stock.entity.do.daily_do import Daily
from module_stock.entity.do.adj_factor_do import AdjFactor
from module_stock.dao.stock_basic_dao import StockBasicDao
from sqlalchemy import func, desc
from sqlalchemy.ext.asyncio import AsyncSession


async def calculate_and_save_stock_changes(*args, **kwargs):
    """
    计算所有股票的3日、5日、10日、20日涨幅，并将结果保存到 stock_basic 表的 other_info 字段中
    考虑前复权、后复权、不复权三种情况
    """
    async for query_db in get_db():
        try:
            # 获取所有股票基本信息
            stock_result = await query_db.execute(select(StockBasic))
            stocks = stock_result.scalars().all()

            updated_count = 0

            for stock in stocks:
                # 获取该股票最近362个交易日的数据（确保足够计算360日涨幅）
                daily_result = await query_db.execute(
                    select(Daily)
                    .where(Daily.ts_code == stock.ts_code)
                    .order_by(desc(Daily.trade_date))
                    .limit(362)
                )
                daily_data = daily_result.scalars().all()

                if len(daily_data) < 3:
                    # 数据不足，跳过该股票
                    continue

                # 获取复权因子数据
                adj_factor_result = await query_db.execute(
                    select(AdjFactor)
                    .where(AdjFactor.ts_code == stock.ts_code)
                    .order_by(desc(AdjFactor.trade_date))
                    .limit(362)
                )
                adj_factor_data = adj_factor_result.scalars().all()

                # 构建复权因子字典，以trade_date为key
                adj_factor_dict = {adj.trade_date.date(): float(adj.adj_factor) for adj in adj_factor_data}

                # 按照交易日期升序排列（从旧到新）
                daily_data_sorted = sorted(daily_data, key=lambda x: x.trade_date)
                close_prices = [float(d.close) for d in daily_data_sorted]
                trade_dates = [d.trade_date.date() for d in daily_data_sorted]

                if len(close_prices) < 3:
                    continue

                # 获取三种价格序列
                raw_prices = close_prices
                qfq_prices = calculate_qfq_prices(close_prices, trade_dates, adj_factor_dict)
                hfq_prices = calculate_hfq_prices(close_prices, trade_dates, adj_factor_dict)

                # 定义需要计算的周期
                periods = [3, 5, 10, 20, 30, 60, 90, 180, 360]

                # 计算三种复权方式下的涨跌幅
                raw_changes = {}
                qfq_changes = {}
                hfq_changes = {}

                for period in periods:
                    raw_changes[f"change_{period}d"] = calculate_change(raw_prices, period) if len(
                        raw_prices) >= period else None
                    qfq_changes[f"change_{period}d"] = calculate_change(qfq_prices, period) if len(
                        qfq_prices) >= period else None
                    hfq_changes[f"change_{period}d"] = calculate_change(hfq_prices, period) if len(
                        hfq_prices) >= period else None

                # 构造要保存的数据
                change_data = {
                    "price_changes": {
                        "raw": raw_changes,
                        "qfq": qfq_changes,  # 前复权
                        "hfq": hfq_changes  # 后复权
                    }
                }

                # 更新 stock_basic 表的 other_info 字段
                update_data = {
                    "id": stock.id,
                    "other_info": json.dumps(change_data, ensure_ascii=False)
                }

                # 创建更新模型
                from module_stock.entity.vo.stock_basic_vo import StockBasicModel
                update_model = StockBasicModel(**update_data)
                update_model.other_info = json.dumps(change_data, ensure_ascii=False)
                # 更新数据库记录
                await StockBasicDao.edit_stock_basic(query_db, update_model, auto_commit=False)
                updated_count += 1

            # 一次性提交所有更改
            await query_db.commit()
            print(f"完成所有股票涨幅数据计算和保存，共更新 {updated_count} 只股票")

            return {"status": "success", "updated_count": updated_count}

        except Exception as e:
            await query_db.rollback()
            print(f"计算和保存股票涨幅时出错: {e}")
            raise e
        finally:
            await query_db.close()


# 计算涨幅的函数 prices价格列表 period计算几日涨幅
def calculate_change(prices, period):
    if len(prices) < period:  # 需要period天的数据来计算period日涨幅
        return None
    # 使用 period 天前的价格作为基准计算涨幅
    base_price = prices[-period]
    current_price = prices[-1]
    if base_price and current_price and base_price != 0:
        return round((current_price - base_price) / base_price * 100, 2)
    return None

# 计算前复权价格  前复权算法：当日收盘价 × 当日复权因子 / 最新复权因子
def calculate_qfq_prices(prices, dates, adj_factor_dict):
    qfq_prices = []
    # 找到最后一个复权因子（最近的日期）
    if not adj_factor_dict:
        error_msg = "Missing adjustment factor in calculate_qfq_prices"
        print(error_msg)
        raise ValueError(error_msg)

    last_adj_factor = None
    for date in reversed(dates):  # 从最近的日期开始找
        if date in adj_factor_dict:
            last_adj_factor = adj_factor_dict[date]
            break

    if not last_adj_factor:
        error_msg = "Missing last_adj_factor in calculate_qfq_prices"
        print(error_msg)
        raise ValueError(error_msg)

    # 计算前复权价格
    for i, date in enumerate(dates):
        if date in adj_factor_dict and adj_factor_dict[date] != 0:
            qfq_price = prices[i] * (adj_factor_dict[date] / last_adj_factor)
            qfq_prices.append(qfq_price)
        else:
            # 如果没有当日复权因子，抛出错误并打印日志
            error_msg = f"Missing adjustment factor for date: {date} in calculate_qfq_prices"
            print(error_msg)
            raise ValueError(error_msg)
    return qfq_prices


# 计算后复权价格 后复权算法：当日收盘价 × 当日复权因子
def calculate_hfq_prices(prices, dates, adj_factor_dict):
    hfq_prices = []
    # 找到第一个复权因子（最早的日期）
    if not adj_factor_dict:
        error_msg = "Missing adjustment factor in calculate_hfq_prices"
        print(error_msg)
        raise ValueError(error_msg)

    first_adj_factor = None
    for date in dates:  # 从最早的日期开始找
        if date in adj_factor_dict:
            first_adj_factor = adj_factor_dict[date]
            break

    if not first_adj_factor:
        error_msg = "Missing first_adj_factor in calculate_hfq_prices"
        print(error_msg)
        raise ValueError(error_msg)

    # 计算后复权价格
    for i, date in enumerate(dates):
        if date in adj_factor_dict and adj_factor_dict[date] != 0:
            hfq_price = prices[i] * adj_factor_dict[date]
            hfq_prices.append(hfq_price)
        else:
            # 如果没有当日复权因子，抛出错误并打印日志
            error_msg = f"Missing adjustment factor for date: {date} in calculate_hfq_prices"
            print(error_msg)
            raise ValueError(error_msg)
    return hfq_prices
